//! 分布式追踪系统测试
//! 
//! 为DistributedTracer trait提供完整的测试和验证

use crate::*;
use std::collections::HashMap;
use std::time::SystemTime;
use std::sync::Arc;
use tokio::sync::RwLock;
use async_trait::async_trait;

/// 高级Mock分布式追踪器
pub struct AdvancedMockDistributedTracer {
    /// 活跃的spans
    active_spans: Arc<RwLock<HashMap<String, TraceSpan>>>,
    /// 已完成的spans
    finished_spans: Arc<RwLock<Vec<TraceSpan>>>,
    /// 统计信息
    statistics: Arc<RwLock<TraceStatistics>>,
    /// 是否启用
    enabled: bool,
    /// 生成trace_id的计数器
    trace_counter: Arc<RwLock<u32>>,
    /// 生成span_id的计数器
    span_counter: Arc<RwLock<u32>>,
}

impl AdvancedMockDistributedTracer {
    pub fn new() -> Self {
        Self {
            active_spans: Arc::new(RwLock::new(HashMap::new())),
            finished_spans: Arc::new(RwLock::new(Vec::new())),
            statistics: Arc::new(RwLock::new(TraceStatistics {
                total_spans: 0,
                active_spans: 0,
                finished_spans: 0,
                avg_span_duration_ms: 0.0,
                sampling_rate: 1.0,
            })),
            enabled: true,
            trace_counter: Arc::new(RwLock::new(0)),
            span_counter: Arc::new(RwLock::new(0)),
        }
    }

    pub fn new_disabled() -> Self {
        let mut tracer = Self::new();
        tracer.enabled = false;
        tracer
    }

    /// 生成唯一的trace_id
    async fn generate_trace_id(&self) -> String {
        let mut counter = self.trace_counter.write().await;
        *counter += 1;
        format!("trace-{:08x}", *counter)
    }

    /// 生成唯一的span_id
    async fn generate_span_id(&self) -> String {
        let mut counter = self.span_counter.write().await;
        *counter += 1;
        format!("span-{:08x}", *counter)
    }

    /// 获取活跃span
    pub async fn get_active_span(&self, span_id: &str) -> Option<TraceSpan> {
        self.active_spans.read().await.get(span_id).cloned()
    }

    /// 获取已完成的spans
    pub async fn get_finished_spans(&self) -> Vec<TraceSpan> {
        self.finished_spans.read().await.clone()
    }

    /// 获取指定trace_id的所有spans
    pub async fn get_spans_by_trace_id(&self, trace_id: &str) -> Vec<TraceSpan> {
        let active = self.active_spans.read().await;
        let finished = self.finished_spans.read().await;
        
        let mut spans = Vec::new();
        
        // 添加活跃spans
        for span in active.values() {
            if span.trace_id == trace_id {
                spans.push(span.clone());
            }
        }
        
        // 添加已完成spans
        for span in finished.iter() {
            if span.trace_id == trace_id {
                spans.push(span.clone());
            }
        }
        
        spans
    }

    /// 清除所有数据
    pub async fn clear_all(&self) {
        self.active_spans.write().await.clear();
        self.finished_spans.write().await.clear();
        *self.statistics.write().await = TraceStatistics {
            total_spans: 0,
            active_spans: 0,
            finished_spans: 0,
            avg_span_duration_ms: 0.0,
            sampling_rate: 1.0,
        };
    }
}

#[async_trait]
impl DistributedTracer for AdvancedMockDistributedTracer {
    async fn start_span(&self, operation_name: &str, parent_span: Option<&TraceSpan>) -> TracingResult<TraceSpan> {
        if !self.enabled {
            return Err(Box::new(TracingError::InvalidOperation("追踪器已禁用".to_string())));
        }

        if operation_name.is_empty() {
            return Err(Box::new(TracingError::InvalidData("操作名称不能为空".to_string())));
        }

        let trace_id = match parent_span {
            Some(parent) => parent.trace_id.clone(),
            None => self.generate_trace_id().await,
        };

        let span_id = self.generate_span_id().await;
        let parent_span_id = parent_span.map(|p| p.span_id.clone());

        let span = TraceSpan {
            trace_id,
            span_id: span_id.clone(),
            parent_span_id,
            operation_name: operation_name.to_string(),
            start_time: SystemTime::now(),
            finish_time: None,
            tags: HashMap::new(),
            logs: Vec::new(),
        };

        // 添加到活跃spans
        self.active_spans.write().await.insert(span_id, span.clone());

        // 更新统计信息
        let mut stats = self.statistics.write().await;
        stats.total_spans += 1;
        stats.active_spans = self.active_spans.read().await.len() as u64;

        Ok(span)
    }

    async fn finish_span(&self, span: &mut TraceSpan) -> TracingResult<()> {
        if !self.enabled {
            return Err(Box::new(TracingError::InvalidOperation("追踪器已禁用".to_string())));
        }

        if span.finish_time.is_some() {
            return Err(Box::new(TracingError::InvalidOperation("Span已经完成".to_string())));
        }

        // 设置完成时间
        span.finish_time = Some(SystemTime::now());

        // 从活跃spans中移除
        let mut active_spans = self.active_spans.write().await;
        active_spans.remove(&span.span_id);

        // 添加到已完成spans
        self.finished_spans.write().await.push(span.clone());

        // 更新统计信息
        let mut stats = self.statistics.write().await;
        stats.finished_spans += 1;
        stats.active_spans = active_spans.len() as u64;

        // 计算平均持续时间
        if let Some(finish_time) = span.finish_time {
            if let Ok(duration) = finish_time.duration_since(span.start_time) {
                let duration_ms = duration.as_millis() as f64;
                if stats.finished_spans == 1 {
                    stats.avg_span_duration_ms = duration_ms;
                } else {
                    let total_duration = stats.avg_span_duration_ms * (stats.finished_spans - 1) as f64 + duration_ms;
                    stats.avg_span_duration_ms = total_duration / stats.finished_spans as f64;
                }
            }
        }

        Ok(())
    }

    async fn add_span_tag(&self, span: &mut TraceSpan, key: &str, value: &str) -> TracingResult<()> {
        if !self.enabled {
            return Err(Box::new(TracingError::InvalidOperation("追踪器已禁用".to_string())));
        }

        if key.is_empty() {
            return Err(Box::new(TracingError::InvalidData("标签键不能为空".to_string())));
        }

        // 添加标签到span
        span.tags.insert(key.to_string(), value.to_string());

        // 更新活跃spans中的标签
        let mut active_spans = self.active_spans.write().await;
        if let Some(active_span) = active_spans.get_mut(&span.span_id) {
            active_span.tags.insert(key.to_string(), value.to_string());
        }

        Ok(())
    }

    async fn add_span_log(&self, span: &mut TraceSpan, message: &str, fields: &HashMap<String, String>) -> TracingResult<()> {
        if !self.enabled {
            return Err(Box::new(TracingError::InvalidOperation("追踪器已禁用".to_string())));
        }

        if message.is_empty() {
            return Err(Box::new(TracingError::InvalidData("日志消息不能为空".to_string())));
        }

        let log = SpanLog {
            timestamp: SystemTime::now(),
            message: message.to_string(),
            fields: fields.clone(),
        };

        // 添加日志到span
        span.logs.push(log.clone());

        // 更新活跃spans中的日志
        let mut active_spans = self.active_spans.write().await;
        if let Some(active_span) = active_spans.get_mut(&span.span_id) {
            active_span.logs.push(log);
        }

        Ok(())
    }

    async fn export_traces(&self, format: TraceExportFormat) -> TracingResult<String> {
        if !self.enabled {
            return Err(Box::new(TracingError::InvalidOperation("追踪器已禁用".to_string())));
        }

        let active_spans = self.active_spans.read().await;
        let finished_spans = self.finished_spans.read().await;

        let mut all_spans = Vec::new();
        all_spans.extend(active_spans.values().cloned());
        all_spans.extend(finished_spans.iter().cloned());

        match format {
            TraceExportFormat::Json => {
                Ok(serde_json::to_string_pretty(&all_spans)
                    .map_err(|e| Box::new(TracingError::InvalidData(format!("JSON序列化失败: {}", e))))?)
            }
            TraceExportFormat::Jaeger => {
                let mut output = String::new();
                output.push_str("# Jaeger Trace Export\n");
                for span in &all_spans {
                    output.push_str(&format!("Trace: {} | Span: {} | Operation: {} | Duration: {:?}\n",
                        span.trace_id,
                        span.span_id,
                        span.operation_name,
                        span.finish_time.and_then(|f| f.duration_since(span.start_time).ok())
                    ));
                    
                    for (key, value) in &span.tags {
                        output.push_str(&format!("  Tag: {} = {}\n", key, value));
                    }
                    
                    for log in &span.logs {
                        output.push_str(&format!("  Log: {} | Fields: {:?}\n", log.message, log.fields));
                    }
                }
                Ok(output)
            }
            TraceExportFormat::OpenTelemetry => {
                let mut output = String::new();
                output.push_str("# OpenTelemetry Trace Export\n");
                for span in &all_spans {
                    output.push_str(&format!("TraceID={} SpanID={} Operation={}\n",
                        span.trace_id, span.span_id, span.operation_name));
                    output.push_str(&format!("StartTime={:?}\n", span.start_time));
                    if let Some(finish_time) = span.finish_time {
                        output.push_str(&format!("EndTime={:?}\n", finish_time));
                    }
                    output.push_str(&format!("Tags={:?}\n", span.tags));
                    output.push_str(&format!("Logs={:?}\n", span.logs));
                    output.push('\n');
                }
                Ok(output)
            }
        }
    }

    async fn get_trace_statistics(&self) -> TracingResult<TraceStatistics> {
        Ok(self.statistics.read().await.clone())
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    // ============ 基础功能测试 ============

    #[tokio::test]
    async fn test_start_span_basic() {
        let tracer = AdvancedMockDistributedTracer::new();

        // 测试根span
        let span = tracer.start_span("test_operation", None).await.unwrap();
        assert_eq!(span.operation_name, "test_operation");
        assert!(span.parent_span_id.is_none());
        assert!(span.trace_id.starts_with("trace-"));
        assert!(span.span_id.starts_with("span-"));
        assert!(span.finish_time.is_none());
        assert!(span.tags.is_empty());
        assert!(span.logs.is_empty());

        // 验证统计信息
        let stats = tracer.get_trace_statistics().await.unwrap();
        assert_eq!(stats.total_spans, 1);
        assert_eq!(stats.active_spans, 1);
        assert_eq!(stats.finished_spans, 0);
    }

    #[tokio::test]
    async fn test_start_child_span() {
        let tracer = AdvancedMockDistributedTracer::new();

        // 创建父span
        let parent_span = tracer.start_span("parent_operation", None).await.unwrap();
        
        // 创建子span
        let child_span = tracer.start_span("child_operation", Some(&parent_span)).await.unwrap();
        
        // 验证子span属性
        assert_eq!(child_span.trace_id, parent_span.trace_id); // 相同的trace_id
        assert_eq!(child_span.parent_span_id, Some(parent_span.span_id.clone()));
        assert_eq!(child_span.operation_name, "child_operation");
        assert_ne!(child_span.span_id, parent_span.span_id); // 不同的span_id

        // 验证统计信息
        let stats = tracer.get_trace_statistics().await.unwrap();
        assert_eq!(stats.total_spans, 2);
        assert_eq!(stats.active_spans, 2);
    }

    #[tokio::test]
    async fn test_start_span_validation() {
        let tracer = AdvancedMockDistributedTracer::new();

        // 测试空操作名称
        let result = tracer.start_span("", None).await;
        assert!(result.is_err());
    }

    #[tokio::test]
    async fn test_finish_span() {
        let tracer = AdvancedMockDistributedTracer::new();

        let mut span = tracer.start_span("test_operation", None).await.unwrap();
        assert!(span.finish_time.is_none());

        // 完成span
        tracer.finish_span(&mut span).await.unwrap();
        assert!(span.finish_time.is_some());

        // 验证统计信息
        let stats = tracer.get_trace_statistics().await.unwrap();
        assert_eq!(stats.finished_spans, 1);
        assert_eq!(stats.active_spans, 0);
        assert!(stats.avg_span_duration_ms >= 0.0);

        // 测试重复完成span
        let result = tracer.finish_span(&mut span).await;
        assert!(result.is_err());
    }

    #[tokio::test]
    async fn test_span_tags() {
        let tracer = AdvancedMockDistributedTracer::new();

        let mut span = tracer.start_span("test_operation", None).await.unwrap();

        // 添加标签
        tracer.add_span_tag(&mut span, "http.method", "GET").await.unwrap();
        tracer.add_span_tag(&mut span, "http.status_code", "200").await.unwrap();
        tracer.add_span_tag(&mut span, "user.id", "12345").await.unwrap();

        // 验证标签
        assert_eq!(span.tags.get("http.method"), Some(&"GET".to_string()));
        assert_eq!(span.tags.get("http.status_code"), Some(&"200".to_string()));
        assert_eq!(span.tags.get("user.id"), Some(&"12345".to_string()));
        assert_eq!(span.tags.len(), 3);

        // 测试覆盖标签
        tracer.add_span_tag(&mut span, "http.method", "POST").await.unwrap();
        assert_eq!(span.tags.get("http.method"), Some(&"POST".to_string()));

        // 测试空键
        let result = tracer.add_span_tag(&mut span, "", "value").await;
        assert!(result.is_err());
    }

    #[tokio::test]
    async fn test_span_logs() {
        let tracer = AdvancedMockDistributedTracer::new();

        let mut span = tracer.start_span("test_operation", None).await.unwrap();

        // 测试简单日志
        tracer.add_span_log(&mut span, "操作开始", &HashMap::new()).await.unwrap();

        // 测试带字段的日志
        let mut fields = HashMap::new();
        fields.insert("user_id".to_string(), "12345".to_string());
        fields.insert("action".to_string(), "login".to_string());
        tracer.add_span_log(&mut span, "用户登录", &fields).await.unwrap();

        // 验证日志
        assert_eq!(span.logs.len(), 2);
        assert_eq!(span.logs[0].message, "操作开始");
        assert!(span.logs[0].fields.is_empty());
        
        assert_eq!(span.logs[1].message, "用户登录");
        assert_eq!(span.logs[1].fields.get("user_id"), Some(&"12345".to_string()));
        assert_eq!(span.logs[1].fields.get("action"), Some(&"login".to_string()));

        // 测试空消息
        let result = tracer.add_span_log(&mut span, "", &HashMap::new()).await;
        assert!(result.is_err());
    }

    #[tokio::test]
    async fn test_export_traces_json() {
        let tracer = AdvancedMockDistributedTracer::new();

        // 创建几个spans
        let mut span1 = tracer.start_span("operation1", None).await.unwrap();
        let mut span2 = tracer.start_span("operation2", Some(&span1)).await.unwrap();
        
        // 添加一些数据
        tracer.add_span_tag(&mut span1, "service", "user-service").await.unwrap();
        tracer.add_span_log(&mut span2, "processing", &HashMap::new()).await.unwrap();
        
        // 完成一个span
        tracer.finish_span(&mut span1).await.unwrap();

        // 导出JSON格式
        let json_export = tracer.export_traces(TraceExportFormat::Json).await.unwrap();
        
        // 验证JSON格式
        let parsed: serde_json::Value = serde_json::from_str(&json_export).unwrap();
        assert!(parsed.is_array());
        
        let spans_array = parsed.as_array().unwrap();
        assert_eq!(spans_array.len(), 2);
    }

    #[tokio::test]
    async fn test_export_traces_jaeger() {
        let tracer = AdvancedMockDistributedTracer::new();

        let mut span = tracer.start_span("test_operation", None).await.unwrap();
        tracer.add_span_tag(&mut span, "component", "database").await.unwrap();
        tracer.finish_span(&mut span).await.unwrap();

        let jaeger_export = tracer.export_traces(TraceExportFormat::Jaeger).await.unwrap();
        
        assert!(jaeger_export.contains("# Jaeger Trace Export"));
        assert!(jaeger_export.contains("test_operation"));
        assert!(jaeger_export.contains("Tag: component = database"));
    }

    #[tokio::test]
    async fn test_export_traces_opentelemetry() {
        let tracer = AdvancedMockDistributedTracer::new();

        let mut span = tracer.start_span("test_operation", None).await.unwrap();
        tracer.add_span_tag(&mut span, "version", "1.0").await.unwrap();
        tracer.finish_span(&mut span).await.unwrap();

        let otel_export = tracer.export_traces(TraceExportFormat::OpenTelemetry).await.unwrap();
        
        assert!(otel_export.contains("# OpenTelemetry Trace Export"));
        assert!(otel_export.contains("TraceID="));
        assert!(otel_export.contains("SpanID="));
        assert!(otel_export.contains("Operation=test_operation"));
    }

    #[tokio::test]
    async fn test_trace_statistics() {
        let tracer = AdvancedMockDistributedTracer::new();

        // 初始统计
        let stats = tracer.get_trace_statistics().await.unwrap();
        assert_eq!(stats.total_spans, 0);
        assert_eq!(stats.active_spans, 0);
        assert_eq!(stats.finished_spans, 0);
        assert_eq!(stats.avg_span_duration_ms, 0.0);

        // 创建spans
        let mut span1 = tracer.start_span("op1", None).await.unwrap();
        let mut span2 = tracer.start_span("op2", None).await.unwrap();
        let mut span3 = tracer.start_span("op3", None).await.unwrap();

        let stats = tracer.get_trace_statistics().await.unwrap();
        assert_eq!(stats.total_spans, 3);
        assert_eq!(stats.active_spans, 3);
        assert_eq!(stats.finished_spans, 0);

        // 完成部分spans
        tracer.finish_span(&mut span1).await.unwrap();
        tracer.finish_span(&mut span2).await.unwrap();

        let stats = tracer.get_trace_statistics().await.unwrap();
        assert_eq!(stats.total_spans, 3);
        assert_eq!(stats.active_spans, 1);
        assert_eq!(stats.finished_spans, 2);
        assert!(stats.avg_span_duration_ms >= 0.0);

        // 完成最后一个span
        tracer.finish_span(&mut span3).await.unwrap();

        let stats = tracer.get_trace_statistics().await.unwrap();
        assert_eq!(stats.active_spans, 0);
        assert_eq!(stats.finished_spans, 3);
    }

    #[tokio::test]
    async fn test_disabled_tracer() {
        let tracer = AdvancedMockDistributedTracer::new_disabled();

        // 所有操作都应该失败
        let result = tracer.start_span("test", None).await;
        assert!(result.is_err());

        let result = tracer.export_traces(TraceExportFormat::Json).await;
        assert!(result.is_err());
    }

    #[tokio::test]
    async fn test_trace_hierarchy() {
        let tracer = AdvancedMockDistributedTracer::new();

        // 创建层次结构：root -> child1 -> grandchild
        //                     -> child2
        let root_span = tracer.start_span("root_operation", None).await.unwrap();
        let child1_span = tracer.start_span("child1_operation", Some(&root_span)).await.unwrap();
        let child2_span = tracer.start_span("child2_operation", Some(&root_span)).await.unwrap();
        let grandchild_span = tracer.start_span("grandchild_operation", Some(&child1_span)).await.unwrap();

        // 验证trace_id一致性
        assert_eq!(child1_span.trace_id, root_span.trace_id);
        assert_eq!(child2_span.trace_id, root_span.trace_id);
        assert_eq!(grandchild_span.trace_id, root_span.trace_id);

        // 验证父子关系
        assert_eq!(child1_span.parent_span_id, Some(root_span.span_id.clone()));
        assert_eq!(child2_span.parent_span_id, Some(root_span.span_id.clone()));
        assert_eq!(grandchild_span.parent_span_id, Some(child1_span.span_id.clone()));

        // 获取同一trace下的所有spans
        let trace_spans = tracer.get_spans_by_trace_id(&root_span.trace_id).await;
        assert_eq!(trace_spans.len(), 4);
    }

    #[tokio::test]
    async fn test_concurrent_spans() {
        let tracer = Arc::new(AdvancedMockDistributedTracer::new());
        let mut handles = vec![];

        // 并发创建spans
        for i in 0..10 {
            let tracer_clone = tracer.clone();
            let handle = tokio::spawn(async move {
                let mut span = tracer_clone.start_span(&format!("concurrent_op_{}", i), None).await.unwrap();
                
                // 添加一些标签和日志
                tracer_clone.add_span_tag(&mut span, "thread_id", &format!("{}", i)).await.unwrap();
                tracer_clone.add_span_log(&mut span, &format!("Processing {}", i), &HashMap::new()).await.unwrap();
                
                // 模拟一些工作
                tokio::time::sleep(tokio::time::Duration::from_millis(10)).await;
                
                tracer_clone.finish_span(&mut span).await.unwrap();
            });
            handles.push(handle);
        }

        // 等待所有任务完成
        for handle in handles {
            handle.await.unwrap();
        }

        // 验证结果
        let stats = tracer.get_trace_statistics().await.unwrap();
        assert_eq!(stats.total_spans, 10);
        assert_eq!(stats.finished_spans, 10);
        assert_eq!(stats.active_spans, 0);

        let finished_spans = tracer.get_finished_spans().await;
        assert_eq!(finished_spans.len(), 10);

        // 验证每个span都有正确的标签
        for span in &finished_spans {
            assert!(span.tags.contains_key("thread_id"));
            assert_eq!(span.logs.len(), 1);
        }
    }

    #[tokio::test]
    async fn test_memory_management() {
        let tracer = AdvancedMockDistributedTracer::new();

        // 创建大量spans
        for i in 0..100 {
            let mut span = tracer.start_span(&format!("operation_{}", i), None).await.unwrap();
            
            // 添加一些数据
            tracer.add_span_tag(&mut span, "iteration", &i.to_string()).await.unwrap();
            tracer.add_span_log(&mut span, &format!("Iteration {}", i), &HashMap::new()).await.unwrap();
            
            // 完成span
            tracer.finish_span(&mut span).await.unwrap();
        }

        // 验证所有spans都被正确处理
        let stats = tracer.get_trace_statistics().await.unwrap();
        assert_eq!(stats.total_spans, 100);
        assert_eq!(stats.finished_spans, 100);
        assert_eq!(stats.active_spans, 0);

        let finished_spans = tracer.get_finished_spans().await;
        assert_eq!(finished_spans.len(), 100);

        // 清理测试
        tracer.clear_all().await;
        let stats = tracer.get_trace_statistics().await.unwrap();
        assert_eq!(stats.total_spans, 0);
        assert_eq!(stats.finished_spans, 0);
    }
}